#!/usr/bin/env python
# -*- encoding:utf-8 -*-

from transformers import BertForSequenceClassification, Trainer, TrainingArguments, BertTokenizerFast
from torch.utils.data import Dataset
import torch
from sklearn.model_selection import train_test_split
import random

id2label = {
    0: "reqAndDuty",
    1: "other",
}

label2id = {
    "reqAndDuty": 0,
    "other": 1,
}

tokenizer = BertTokenizerFast.from_pretrained('hfl/chinese-bert-wwm')
model = BertForSequenceClassification.from_pretrained(
    'hfl/chinese-bert-wwm', num_labels=2, id2label=id2label, label2id=label2id)
if torch.cuda.is_available():
    model = model.to("cuda")
else:
    model = model.to("cpu")

# layer.8,layer.9,layer.10,layer.11 变更权重
for (name, param) in model.bert.named_parameters():
    param.requires_grad = False
    if "layer.8." in name or "layer.9." in name or "layer.10." in name or "layer.11." in name:
        param.requires_grad = True


class MyDataset(Dataset):
    def __init__(self, itmes, tokenizer, max_len):
        self.itmes = itmes
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        return len(self.itmes)

    def __getitem__(self, idx):
        (text, label) = self.itmes[idx]
        encoding = self.tokenizer(
            text,
            add_special_tokens=True,
            max_length=self.max_len,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )
        return {
            "input_ids": encoding["input_ids"].squeeze(0),
            "attention_mask": encoding["attention_mask"].squeeze(0),
            "label": torch.tensor(label, dtype=torch.long),
        }


def loadAllData():
    itemList = []
    for (fileName, label) in [("../data/lines_req.txt.data", 0), ("../data/lines_duty.txt.data", 0), ("../data/lines_other.txt.data", 1)]:
        with open(fileName, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                itemList.append((line, label))

    random.shuffle(itemList)
    return itemList


# 加载数据集合
dataItemList = loadAllData()

trainItems, testItems = train_test_split(dataItemList, test_size=0.2)

train_dataset = MyDataset(trainItems, tokenizer, 100)
test_dataset = MyDataset(testItems, tokenizer, 100)

training_args = TrainingArguments(
    output_dir='./results',          # output directory
    num_train_epochs=3,              # total # of training epochs
    per_device_train_batch_size=1024,  # batch size per device during training
    per_device_eval_batch_size=64,   # batch size for evaluation
    warmup_steps=500,                # number of warmup steps for learning rate scheduler
    weight_decay=0.01,               # strength of weight decay
    logging_dir='./logs',            # directory for storing logs
    eval_strategy="epoch",
)

trainer = Trainer(
    model=model,
    args=training_args,                  # training arguments, defined above
    train_dataset=train_dataset,         # training dataset
    eval_dataset=test_dataset            # evaluation dataset
)

trainer.train()
